#include "cim_tools.h"
#include <cassert>
#include <cstdio>
#include <stdlib.h>
#include <unistd.h>
#include <stdio.h>
#include <fcntl.h>
#include <vector>
#include <iostream>
#include <fstream>
#include <cinttypes>
#include <filesystem>
#include <fcntl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include "ai_common.h"
#include <random>




namespace cim_tools {
    void sw_mmult(
        const MatrixS8 input,
        const MatrixS8 weight,
        MatrixS8 output,
        uint8_t shift_num,
        const char * _save_path
    ){
        // int i, j, k;
        MatrixS32 golden_out_data_full;
        MatrixS32 golden_out_data_shift;

        golden_out_data_full.alloc(weight.H, input.W);
        golden_out_data_shift.alloc(weight.H, input.W);
        assert(weight.W == input.H && "Weight width must match input height");
        // Matrix multiplication: golden_out_data_full = weight * input
        for (size_t i = 0; i < weight.H; ++i) {
            for (size_t j = 0; j < input.W; ++j) {
                int32_t sum = 0;
                for (size_t k = 0; k < weight.W; ++k) {
                    int8_t w = weight.data[i * weight.W + k];
                    int8_t in = input.data[k * input.W + j];
                    sum += static_cast<int32_t>(w) * static_cast<int32_t>(in);
                }
                golden_out_data_full.data[i * input.W + j] = sum;
            }
        }

        // Right shift by shift_num to get golden_out_data_shift
        for (size_t i = 0; i < weight.H; ++i) {
            for (size_t j = 0; j < input.W; ++j) {
                int32_t val = golden_out_data_full.data[i * input.W + j];
                golden_out_data_shift.data[i * input.W + j] = val >> shift_num;
            }
        }

        // Truncate to 8 bits and store in output
        output.alloc(weight.H, input.W);
        for (size_t i = 0; i < weight.H; ++i) {
            for (size_t j = 0; j < input.W; ++j) {
                int32_t val = golden_out_data_shift.data[i * input.W + j];
                // Clamp to int8_t range
                if (val > 127) val = 127;
                if (val < -128) val = -128;
                // output.data[i * input.W + j] = static_cast<int8_t>(val);
                output.data[j * weight.H + i] = static_cast<int8_t>(val);
            }
        }

        char filename[1024];
        // log results
        check_folder(_save_path);
        sprintf(filename, "%s/output.bin", _save_path);
        write_bin_file(filename, output.data, output.numel());
        sprintf(filename, "%s/input.bin", _save_path);
        write_bin_file(filename, input.data, input.numel());
        sprintf(filename, "%s/weight.bin", _save_path);
        write_bin_file(filename, weight.data, weight.numel());
    }

    void inst_gen_weight_dma(std::vector<uint32_t> *inst_buf, uint64_t axi_saddr, int num) {
        int reverved_transfer_num = num;
        unsigned int sram_saddr = 1<<19;
        while (reverved_transfer_num > 0){
            int transfer_num;
            if(reverved_transfer_num > 256){
                transfer_num = 256;
            } else {
                transfer_num = reverved_transfer_num;
            }
            uint32_t inst31_0 =  (( axi_saddr & 0x1FFFF )<< 15)| transfer_num; //{axi_saddr[16:0] , transfer_wt_num[14:0]}
            uint32_t inst63_32 = ((sram_saddr & 0x1FF)<<23) | (axi_saddr>>17); //{sram_saddr[8:0],axi_saddr[39:17] 23}
            uint32_t inst95_64 = (1<<11) |  (sram_saddr>>9) ;//{sram_saddr[19:9]}
            uint32_t inst127_96 = 1<<28;

            inst_buf->push_back(inst31_0);
            inst_buf->push_back(inst63_32);
            inst_buf->push_back(inst95_64);
            inst_buf->push_back(inst127_96);

            sram_saddr = sram_saddr + 256;
            axi_saddr = axi_saddr + 0x1000;
            reverved_transfer_num = reverved_transfer_num - 256;
        }
    }

    void inst_gen_input_dma(std::vector<uint32_t> *inst_buf,  uint64_t axi_saddr, int num) {
        int reverved_transfer_num = num;
        unsigned int sram_saddr = 0;
        while (reverved_transfer_num > 0){
            int transfer_num;
            if(reverved_transfer_num > 256){
                transfer_num = 256;
            } else {
                transfer_num = reverved_transfer_num;
            }
            uint32_t inst31_0 =  (( axi_saddr & 0x1FFFF )<< 15)| transfer_num; //{axi_saddr[16:0] , transfer_wt_num[14:0]}
            uint32_t inst63_32 = ((sram_saddr & 0x1FF)<<23) | (axi_saddr>>17); //{sram_saddr[8:0],axi_saddr[39:17] 23}
            uint32_t inst95_64 = (1<<11) |  (sram_saddr>>9) ;//{sram_saddr[19:9]}
            uint32_t inst127_96 = 2<<28;

            inst_buf->push_back(inst31_0);
            inst_buf->push_back(inst63_32);
            inst_buf->push_back(inst95_64);
            inst_buf->push_back(inst127_96);

            sram_saddr = sram_saddr + 256;
            axi_saddr = axi_saddr + 0x1000;
            reverved_transfer_num = reverved_transfer_num - 256;
        }
    }

    void inst_gen_output_dma(std::vector<uint32_t> *inst_buf,  uint64_t axi_saddr, int num) {
        int reverved_transfer_num = num;
        unsigned int sram_saddr = 1<<18; 
        while (reverved_transfer_num > 0){
            int transfer_num;
            if(reverved_transfer_num > 256){
                transfer_num = 256;
            } else {
                transfer_num = reverved_transfer_num;
            }
            uint32_t inst31_0 =  (( axi_saddr & 0x1FFFF )<< 15)| transfer_num; //{axi_saddr[16:0] , transfer_wt_num[14:0]}
            uint32_t inst63_32 = ((sram_saddr & 0x1FF)<<23) | (axi_saddr>>17); //{sram_saddr[8:0],axi_saddr[39:17] 23}
            uint32_t inst95_64 = (0<<11) |  (sram_saddr>>9) ;//{sram_saddr[19:9]}
            uint32_t inst127_96 = 2<<28;

            inst_buf->push_back(inst31_0);
            inst_buf->push_back(inst63_32);
            inst_buf->push_back(inst95_64);
            inst_buf->push_back(inst127_96);

            sram_saddr = sram_saddr + 256;
            axi_saddr = axi_saddr + 0x1000;
            reverved_transfer_num = reverved_transfer_num - 256;
        }
    }

    void inst_gen_finish(std::vector<uint32_t> *inst_buf){
        uint32_t inst31_0   = 0;
        uint32_t inst63_32  = 0;
        uint32_t inst95_64  = 0;
        uint32_t inst127_96 = 4<<28;
        inst_buf->push_back(inst31_0);
        inst_buf->push_back(inst63_32);
        inst_buf->push_back(inst95_64);
        inst_buf->push_back(inst127_96); 
    }

    void inst_gen_param(std::vector<uint32_t> *inst_buf, uint8_t cfg_cim_relu_enable, uint8_t cim_shift_num, uint8_t cim_in_num){
        uint32_t inst31_0    = (cfg_cim_relu_enable<<16) | (cim_shift_num<<8) | cim_in_num ;
        uint32_t inst63_32   = 0;
        uint32_t inst95_64   = 0;
        uint32_t inst127_96  = 0;

        inst_buf->push_back(inst31_0);
        inst_buf->push_back(inst63_32);
        inst_buf->push_back(inst95_64);
        inst_buf->push_back(inst127_96);
    }

    void inst_gen_cim(std::vector<uint32_t> *inst_buf){
        uint32_t inst31_0   = 0;
        uint32_t inst63_32  = 0;
        uint32_t inst95_64  = 0;
        uint32_t inst127_96 = 3<<28;
        inst_buf->push_back(inst31_0);
        inst_buf->push_back(inst63_32);
        inst_buf->push_back(inst95_64);
        inst_buf->push_back(inst127_96);
    }
    static void MatrixS8_randinit_input(MatrixS8 matrix){
        assert(matrix.numel() > 0);
        static bool seeded = true;
        if (seeded) {
            srand((unsigned int)time(nullptr));
        }
        for (size_t i = 0; i < matrix.numel(); i++){
            matrix.data[i] = rand();
        }
    }

    static void MatrixS8_randinit_weight(MatrixS8 matrix){
        assert(matrix.numel() > 0);
        static bool seeded = true;
        if (seeded) {
            srand((unsigned int)time(nullptr));
        }
        for (size_t i = 0; i < matrix.numel(); i++){
            matrix.data[i] = rand();
        }
    }
    void LayerConfig::data_rand_generator() {
        assert(input.data != nullptr && "Input matrix must be allocated");
        assert(weight.data != nullptr && "Weight matrix must be allocated");
        MatrixS8_randinit_input(input);
        MatrixS8_randinit_weight(weight);
    }

    static int open_checked(const char* name, int flags) {
        int fd = open(name, flags);
        if (fd == -1) {
            printf("[error] when open `%s`: %s\n", name, strerror(errno));
        }
        return fd;
    }

    void LayerConfig::wt_inst_gen(){
        assert(this->input.H > 0); // 使用this指针访问当前对象的成员
        this->inst_wt.clear();
        inst_gen_param(&this->inst_wt, 0, this->shift_num, this->input.W);
        inst_gen_weight_dma(&this->inst_wt, this->BASE_WT_ADDR, this->cfg_cim_wet_num);
        inst_gen_finish(&this->inst_wt);
        exportVecToBin(this->inst_wt, "inst_wt.bin");
    }


    void LayerConfig::run_inst_gen(){
        assert(this->input.H > 0);
        this->inst_run.clear();
        inst_gen_param(&this->inst_run, 0, this->shift_num, this->input.W);
        // inst_gen_weight_dma(&this->inst_wt, this->BASE_WT_ADDR, this->cfg_cim_wet_num);
        // inst_gen_input_dma(&this->inst_run, this->BASE_IN_ADDR, this->input.H * this->input.W * 8 / 128);
        inst_gen_input_dma(&this->inst_run, this->BASE_IN_ADDR, 128);
        inst_gen_cim(&this->inst_run);
        // inst_gen_output_dma(&this->inst_run, this->BASE_RS_ADDR, this->output.H * this->output.W * 8 / 128);
        inst_gen_output_dma(&this->inst_run, this->BASE_RS_ADDR, 64);
        inst_gen_finish(&this->inst_run);
        exportVecToBin(this->inst_run, "inst_run.bin");
    }

    void LayerConfig::single_run_inst_gen(){
        this->inst_singlerun.clear();
        this->inst_input_dma.clear();
        this->inst_output_dma.clear();

        inst_gen_param(&this->inst_input_dma, 0, this->shift_num, this->input.W);
        inst_gen_input_dma(&this->inst_input_dma, this->BASE_IN_ADDR, 128);
        inst_gen_finish(&this->inst_input_dma);

        inst_gen_cim(&this->inst_singlerun);
        inst_gen_finish(&this->inst_singlerun);

        inst_gen_output_dma(&this->inst_output_dma, this->BASE_RS_ADDR, 64);
        inst_gen_finish(&this->inst_output_dma);

        exportVecToBin(this->inst_input_dma, "singleinst_inputdma.bin");
        exportVecToBin(this->inst_singlerun, "singleinst_run.bin");
        exportVecToBin(this->inst_output_dma, "singleinst_outputdma.bin");
        
    }

    // Offline model 
    void CIM_offline_utils::run(const char * _save_path) {
        assert(layer_config.input.data != nullptr && "Input matrix must be allocated");
        assert(layer_config.weight.data != nullptr && "Weight matrix must be allocated");
        assert(layer_config.input.H % 128 == 0 && "Input height must be a multiple of 128");
        assert(layer_config.input.W < 64 && "Input width exceeds maximum output number 64");
        assert(layer_config.weight.H % 16 == 0 && "Weight height must be a multiple of 16");
        assert(layer_config.weight.W % 128 == 0 && "Weight width must be a multiple of 128");
        layer_config.output.alloc(layer_config.weight.H, layer_config.input.W);    
        // Perform software matrix multiplication
        sw_mmult(
            layer_config.input,
            layer_config.weight,
            layer_config.output,
            layer_config.shift_num,
            _save_path
        );
    }

    // Online model
    #define SRAM_BASE   0x640000000ULL
    #define SRAM_SIZE   (16 * 1024)
    #define DDR_BASE    0x95000000ULL
    #define DDR_SIZE    (256 * 1024 * 1024)
    #define APB_BASE    0x37000000ULL

    static void* pmem_mmap(int devmem_fd, size_t base, size_t length) {
        long page_size = sysconf(_SC_PAGESIZE);
        assert(base % page_size == 0); // 页对齐检查
        void* ptr = mmap(
            NULL,
            length,
            PROT_READ | PROT_WRITE,
            MAP_SHARED,
            devmem_fd,
            base
        );
        if (ptr == MAP_FAILED) {
            printf("[error] mmap /dev/mem at %p: %s\n", (void*)base, strerror(errno));
            exit(1);
        }
        return ptr;
    }

    void CIM_online_utils::pmem_mmap_init() {
        devmem_fd = open_checked("/dev/mem", O_SYNC | O_RDWR);
        dma_inst = (uint32_t*)pmem_mmap(devmem_fd, SRAM_BASE, SRAM_SIZE);
        ddrbuf   = (uint32_t*)pmem_mmap(devmem_fd, DDR_BASE, DDR_SIZE); 
        apb_regs = (apb_regs_block*)pmem_mmap(devmem_fd, APB_BASE, sizeof(apb_regs_block));
        // 其他寄存器空间可按需映射
        apb_regs->TOP_AI_CHOOSE = 0x3;
    }

    void CIM_online_utils::wt_init(){
        memcpy(const_cast<uint32_t*>(dma_inst), CIM_online_utils::layer_config.inst_wt.data(), CIM_online_utils::layer_config.inst_wt.size() * sizeof(uint32_t)); 
        // COPY WEIGHT
        uint64_t wt_offset = (layer_config.BASE_WT_ADDR - CIM_online_utils::DDR_BASE_ADDR) / 4;
        CIM_online_utils::wet_transfer(this->layer_config.weight, this->layer_config._weight);
        printf("layer_config._weight size: %d\n", this->layer_config._weight.numel());
        memcpy(const_cast<uint32_t*>(ddrbuf) + wt_offset, this->layer_config._weight.data, this->layer_config._weight.numel() * sizeof(uint32_t));
        printf("Copy weight to DDR successfully. \n");
        // apb_regs->CIM_CFG_FINISH_CLR = 1;
        // apb_regs->CIM_CFG_FINISH_CLR = 0;
        apb_regs->CIM_TOP_START = 0;
        apb_regs->CIM_TOP_START = 1;


        while(apb_regs->CIM_CFG_TOP_FINISH == 0){
        };
        printf("Write weight successfully. \n");
        apb_regs->CIM_TOP_START = 0;
        apb_regs->CIM_CFG_FINISH_CLR = 1;
        apb_regs->CIM_CFG_FINISH_CLR = 0;
    }

    void CIM_online_utils::cfg_init(){
        apb_regs->CIM_ROW_CONFIG_BITS_0 = 0x33010;
        apb_regs->CIM_ROW_CONFIG_BITS_1  = 0x0;
        apb_regs->CIM_ROW_CONFIG_BITS_2  = 0x0;
        apb_regs->CIM_ROW_CONFIG_BITS_3  = 0x0;
        apb_regs->CIM_ROW_CONFIG_BITS_4  = 0x0;
        apb_regs->CIM_ROW_CONFIG_BITS_5  = 0x0;
        apb_regs->CIM_ROW_CONFIG_BITS_6  = 0x0;
        apb_regs->CIM_ROW_CONFIG_BITS_7  = 0x0;
        apb_regs->CIM_ROW_CONFIG_BITS_8  = 0x0;
        apb_regs->CIM_ROW_CONFIG_BITS_9  = 0x0;
        apb_regs->CIM_ROW_CONFIG_BITS_10 = 0x0;
        apb_regs->CIM_ROW_CONFIG_BITS_11 = 0x0;
        apb_regs->CIM_ROW_CONFIG_BITS_12 = 0x0;
        apb_regs->CIM_ROW_CONFIG_BITS_13 = 0x0;
        apb_regs->CIM_ROW_CONFIG_BITS_14 = 0x0;

        // apb_regs->CIM_ROW_CONFIG_BITS_1 = 0x33010;
        // apb_regs->CIM_ROW_CONFIG_BITS_2 = 0x33010;
        // apb_regs->CIM_ROW_CONFIG_BITS_3 = 0x33010;
        // apb_regs->CIM_ROW_CONFIG_BITS_4 = 0x33010;
        // apb_regs->CIM_ROW_CONFIG_BITS_5 = 0x33010;
        // apb_regs->CIM_ROW_CONFIG_BITS_6 = 0x33010;
        // apb_regs->CIM_ROW_CONFIG_BITS_7 = 0x33010;
        // apb_regs->CIM_ROW_CONFIG_BITS_8 = 0x33010;
        // apb_regs->CIM_ROW_CONFIG_BITS_9 = 0x33010;
        // apb_regs->CIM_ROW_CONFIG_BITS_10 = 0x33010;
        // apb_regs->CIM_ROW_CONFIG_BITS_11 = 0x33010;
        // apb_regs->CIM_ROW_CONFIG_BITS_12 = 0x33010;
        // apb_regs->CIM_ROW_CONFIG_BITS_13 = 0x33010;
        // apb_regs->CIM_ROW_CONFIG_BITS_14 = 0x33010;

        // apb_regs->CIM_CFG_CIM_PE_ENABLE0_31 = 0x0; // Enable all PEs
        // apb_regs->CIM_CFG_CIM_PE_ENABLE32_63 = 0x0; // Enable all PEs
        apb_regs->CIM_CFG_CIM_PE_ENABLE0_31 = 0xFFFFFFFF; // Enable all PEs
        apb_regs->CIM_CFG_CIM_PE_ENABLE32_63 = 0xFFFFFFFF; // Enable all PEs
    }

    void CIM_online_utils::act_transfer(const MatrixS8 input, VectorUS32 &_input) {
        assert(input.data != nullptr && "Input matrix must be allocated");
        assert((input.H * input.W ) % 4 == 0 && "Input size must be a multiple of 4");
        _input.alloc(input.H * input.W / 4);
        uint32_t offset = 0;
        for (size_t i = 0; i < input.W; i++) {
            for (int m = 0; m < 8; m++){
                for (int j = 0; j < CIM_COL; j++){
                    uint32_t ddr32[4] = {0, 0, 0, 0};
                    for (int b = 0; b < 128; b++){
                        int index = (int) (b / 32);
                        ddr32[index] = ddr32[index]  | (( ((uint8_t)(input.data[(j*128+b)*input.W + i]>>m)) & 0x1 ) << (b-32*index) );
                    }
                    _input.data[offset+0] = ddr32[0];
                    _input.data[offset+1] = ddr32[1];
                    _input.data[offset+2] = ddr32[2];
                    _input.data[offset+3] = ddr32[3];
                    offset = offset + 4;
                }
            }
        }
    }

    void CIM_online_utils::wet_transfer(const MatrixS8 weight, VectorUS32 & _weight) {
        assert(weight.data != nullptr && "Weight matrix must be allocated");
        assert((weight.H * weight.W ) % 4 == 0 && "Weight size must be a multiple of 4");
        _weight.alloc(weight.H * weight.W / 4);
        printf("_weight.size:%d\n", _weight.numel());
        uint32_t offset = 0;
        
        for (int i = 0; i < CIM_ROW; i++) {
            for (int j = 0; j < CIM_COL; j++){
                for(int m = 0; m < 16; m++) {
                    for (int b = 0; b < 8; b++){    
                        uint32_t ddr32[4] = {0, 0, 0, 0};           
                        for (int k = 0; k < 128; k++){
                            int index = (int) (k / 32);
                            ddr32[index] = ddr32[index]  | (( ((uint8_t)(weight.data[(i*16+m)*128*4 + (128*j+k)]>>b)) & 0x1 ) << (k-32*index) );
                        }
                        _weight.data[offset+0] = ddr32[0];
                        _weight.data[offset+1] = ddr32[1];
                        _weight.data[offset+2] = ddr32[2];
                        _weight.data[offset+3] = ddr32[3];
                        offset = offset + 4;
                    }
                }
            }
        }

    }

    void CIM_online_utils::run(const char * _save_path){
        Timer timer_input_cpy, timer_compute, timer_output_cpy;
        timer_input_cpy.reset();
        timer_compute.reset();
        timer_output_cpy.reset();
        int times = 1;
        for (int t = 0; t < times; t++){
            timer_input_cpy.start();
            memcpy(const_cast<uint32_t*>(dma_inst), CIM_online_utils::layer_config.inst_run.data(), CIM_online_utils::layer_config.inst_run.size() * sizeof(uint32_t)); 
            // COPY INPUT 
            uint64_t in_offset = (layer_config.BASE_IN_ADDR - CIM_online_utils::DDR_BASE_ADDR) / 4;
            this->act_transfer(this->layer_config.input, this->layer_config._input);
            memcpy(const_cast<uint32_t*>(ddrbuf) + in_offset, layer_config._input.data, layer_config._input.numel() * sizeof(uint32_t));
            timer_input_cpy.stop();
            apb_regs->CIM_TOP_START = 0;
            apb_regs->CIM_CFG_FINISH_CLR = 1;
            apb_regs->CIM_CFG_FINISH_CLR = 0;
            timer_compute.start();
            apb_regs->CIM_TOP_START = 1;
            while(apb_regs->CIM_CFG_TOP_FINISH == 0);
            timer_compute.stop();
            apb_regs->CIM_TOP_START = 0;
            apb_regs->CIM_CFG_FINISH_CLR = 1;
            apb_regs->CIM_CFG_FINISH_CLR = 0;
            
            timer_output_cpy.start();
            //COPY OUTPUT
            uint64_t rs_offset = (layer_config.BASE_RS_ADDR - CIM_online_utils::DDR_BASE_ADDR) / 4;
            // printf("BASE_RS_ADDR: 0x%lx \n", layer_config.BASE_RS_ADDR);
            // printf("DDR_BASE_ADDR: 0x%lx \n", CIM_online_utils::DDR_BASE_ADDR);
            // printf("Result numel: %zu \n", layer_config.output.numel());
            memcpy(layer_config.output.data, const_cast<uint32_t*>(ddrbuf) + rs_offset, layer_config.output.numel() * sizeof(int8_t));
            timer_output_cpy.stop();
        }
        // Print timing information
        timer_input_cpy.print("Input Copy");
        timer_compute.print("Compute");
        timer_output_cpy.print("Output Copy");

        float flops = (float)layer_config.input.H * layer_config.input.W * layer_config.weight.H * layer_config.weight.W * 2 * times;
        float time = (float)timer_compute.time / 1000000.0f; // Convert microseconds to seconds
        float gflops = flops / (time * 1e9); // Convert to GFLOPS
        printf("GFLOPS: %.2f\n", gflops);
        // Save output
        char filename[1024];
        // log results      
        // char _save_path[128] = "./runfile_log";
        check_folder(_save_path);
        sprintf(filename, "%s/output.bin", _save_path);
        write_bin_file(filename, layer_config.output.data, layer_config.output.numel());
        sprintf(filename, "%s/input.bin", _save_path);
        write_bin_file(filename, layer_config._input.data, layer_config._input.numel());
        sprintf(filename, "%s/weight.bin", _save_path); 
        write_bin_file(filename, layer_config._weight.data, layer_config._weight.numel());
    }

    void CIM_online_utils::step_run(const char * _save_path){
        Timer timer_input_cpy, timer_compute, timer_output_cpy;
        timer_input_cpy.reset();
        timer_compute.reset();
        timer_output_cpy.reset();

        timer_input_cpy.start();
        memcpy(const_cast<uint32_t*>(dma_inst), CIM_online_utils::layer_config.inst_input_dma.data(), CIM_online_utils::layer_config.inst_input_dma.size() * sizeof(uint32_t)); 
        // COPY INPUT 
        uint64_t in_offset = (layer_config.BASE_IN_ADDR - CIM_online_utils::DDR_BASE_ADDR) / 4;
        this->act_transfer(this->layer_config.input, this->layer_config._input);
        memcpy(const_cast<uint32_t*>(ddrbuf) + in_offset, layer_config._input.data, layer_config._input.numel() * sizeof(uint32_t));

        apb_regs->CIM_TOP_START = 0;
        apb_regs->CIM_TOP_START = 1;
        printf("Start CIM CPY Input. \n");
        while(apb_regs->CIM_CFG_TOP_FINISH == 0){
            // usleep(1);
        };
        printf("CIM CPY Input successfully. \n");
        apb_regs->CIM_TOP_START = 0;
        apb_regs->CIM_CFG_FINISH_CLR = 1;
        apb_regs->CIM_CFG_FINISH_CLR = 0;
        timer_input_cpy.stop();


        memcpy(const_cast<uint32_t*>(dma_inst), CIM_online_utils::layer_config.inst_singlerun.data(), CIM_online_utils::layer_config.inst_singlerun.size() * sizeof(uint32_t)); 
        
        printf("Start CIM Compute. \n");
        timer_compute.start();
        apb_regs->CIM_CFG_FINISH_CLR = 1;
        apb_regs->CIM_CFG_FINISH_CLR = 0;
        apb_regs->CIM_TOP_START = 1;
        while(apb_regs->CIM_CFG_TOP_FINISH == 0);
        timer_compute.stop();
        printf("CIM Compute successfully. \n");
        
        apb_regs->CIM_TOP_START = 0;
        apb_regs->CIM_CFG_FINISH_CLR = 1;
        apb_regs->CIM_CFG_FINISH_CLR = 0;

        timer_output_cpy.start();
        memcpy(const_cast<uint32_t*>(dma_inst), CIM_online_utils::layer_config.inst_output_dma.data(), CIM_online_utils::layer_config.inst_output_dma.size() * sizeof(uint32_t)); 

        apb_regs->CIM_TOP_START = 1;
        printf("Start RS COPY. \n");
        while(apb_regs->CIM_CFG_TOP_FINISH == 0){
            // usleep(1);
        };
        printf("RS COPY successfully. \n");
        apb_regs->CIM_TOP_START = 0;
        apb_regs->CIM_CFG_FINISH_CLR = 1;
        apb_regs->CIM_CFG_FINISH_CLR = 0;

        //COPY OUTPUT
        uint64_t rs_offset = (layer_config.BASE_RS_ADDR - CIM_online_utils::DDR_BASE_ADDR) / 4;
        memcpy(layer_config.output.data, const_cast<uint32_t*>(ddrbuf) + rs_offset, layer_config.output.numel() * sizeof(int8_t));
        timer_output_cpy.stop();
        // Print timing information
        timer_input_cpy.print("Input Copy");
        timer_compute.print("Compute");
        timer_output_cpy.print("Output Copy");

        float flops = (float)layer_config.input.H * layer_config.input.W * layer_config.weight.H * layer_config.weight.W * 2;
        float time = (float)timer_compute.time / 1000000.0f; // Convert microseconds to seconds
        float gflops = flops / (time * 1e9); // Convert to GFLOPS
        printf("GFLOPS: %.2f\n", gflops);
        // Save output
        char filename[1024];
        // log results      
        // char _save_path[128] = "./runfile_log";
        check_folder(_save_path);
        sprintf(filename, "%s/output.bin", _save_path);
        write_bin_file(filename, layer_config.output.data, layer_config.output.numel());
        sprintf(filename, "%s/input.bin", _save_path);
        write_bin_file(filename, layer_config._input.data, layer_config._input.numel());
        sprintf(filename, "%s/weight.bin", _save_path); 
        write_bin_file(filename, layer_config._weight.data, layer_config._weight.numel());
    }    

}